#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
性能对比测试
比较原始方式 vs 队列模式的处理性能
"""

import os
import time
import tempfile
import shutil
from tools.FileSearchTools import FileSearchTool
from tools.QueuedFileSearchTools import QueuedFileSearchTool


def create_test_files(base_dir: str, file_count: int) -> str:
    """创建大量测试文件"""
    test_dir = os.path.join(base_dir, f"perf_test_{file_count}")
    os.makedirs(test_dir, exist_ok=True)
    
    print(f"正在创建 {file_count} 个测试文件...")
    
    for i in range(file_count):
        # 创建子目录结构
        if i % 100 == 0:
            sub_dir = os.path.join(test_dir, f"batch_{i//100}")
            os.makedirs(sub_dir, exist_ok=True)
        else:
            sub_dir = test_dir
        
        # 创建不同类型的文件
        extensions = ['.txt', '.py', '.json', '.md', '.csv', '.log']
        ext = extensions[i % len(extensions)]
        
        file_path = os.path.join(sub_dir, f"test_file_{i:06d}{ext}")
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(f"测试文件 {i}\n类型: {ext}\n创建时间: {time.ctime()}\n")
    
    print(f"✅ 测试文件创建完成: {test_dir}")
    return test_dir


def test_original_method(test_dir: str, file_count: int):
    """测试原始方法"""
    print(f"\n🔄 测试原始方法 (文件数: {file_count})")
    print("-" * 40)
    
    try:
        start_time = time.time()
        
        # 使用原始FileSearchTool
        search_tool = FileSearchTool(collection_name="original_perf_test")
        search_tool.clear_collection()
        
        # 扫描文件
        scan_start = time.time()
        file_list = search_tool.scan_directory(test_dir, max_files=-1)
        scan_time = time.time() - scan_start
        
        # 索引文件
        index_start = time.time()
        search_tool.index_files(file_list, batch_size=50)
        index_time = time.time() - index_start
        
        total_time = time.time() - start_time
        
        # 统计结果
        stats = search_tool.get_collection_stats()
        
        print(f"✅ 原始方法完成")
        print(f"   扫描时间: {scan_time:.2f}秒")
        print(f"   索引时间: {index_time:.2f}秒")
        print(f"   总时间: {total_time:.2f}秒")
        print(f"   索引文件数: {stats['total_entities']}")
        print(f"   平均速度: {stats['total_entities']/total_time:.1f} 文件/秒")
        
        return {
            'method': 'original',
            'scan_time': scan_time,
            'index_time': index_time,
            'total_time': total_time,
            'indexed_count': stats['total_entities'],
            'avg_speed': stats['total_entities']/total_time
        }
        
    except Exception as e:
        print(f"❌ 原始方法测试失败: {e}")
        return None


def test_queued_method(test_dir: str, file_count: int):
    """测试队列方法"""
    print(f"\n⚡ 测试队列方法 (文件数: {file_count})")
    print("-" * 40)
    
    try:
        start_time = time.time()
        
        # 使用队列模式FileSearchTool
        search_tool = QueuedFileSearchTool(
            collection_name="queued_perf_test",
            queue_size=5000,
            batch_size=100,
            num_workers=3
        )
        search_tool.clear_collection()
        
        # 启动队列索引
        search_tool.start_queued_indexing(test_dir, max_files=-1)
        
        # 等待完成
        search_tool.wait_completion(timeout=600)  # 10分钟超时
        
        total_time = time.time() - start_time
        
        # 获取最终统计
        progress = search_tool.get_progress()
        stats = search_tool.get_collection_stats()
        
        print(f"✅ 队列方法完成")
        print(f"   总时间: {total_time:.2f}秒")
        print(f"   扫描文件数: {progress.scanned_files}")
        print(f"   索引文件数: {progress.indexed_files}")
        print(f"   失败文件数: {progress.failed_files}")
        print(f"   数据库记录: {stats['total_entities']}")
        print(f"   扫描速度: {progress.scan_speed:.1f} 文件/秒")
        print(f"   索引速度: {progress.index_speed:.1f} 文件/秒")
        
        return {
            'method': 'queued',
            'total_time': total_time,
            'scanned_count': progress.scanned_files,
            'indexed_count': progress.indexed_files,
            'failed_count': progress.failed_files,
            'db_records': stats['total_entities'],
            'scan_speed': progress.scan_speed,
            'index_speed': progress.index_speed
        }
        
    except Exception as e:
        print(f"❌ 队列方法测试失败: {e}")
        return None


def compare_performance():
    """性能对比测试"""
    print("🏁 文件处理性能对比测试")
    print("=" * 50)
    
    # 测试不同文件数量
    test_sizes = [1000, 5000, 10000]
    
    results = []
    temp_base = tempfile.mkdtemp(prefix="perf_test_")
    
    try:
        for size in test_sizes:
            print(f"\n📊 测试文件数量: {size}")
            print("=" * 30)
            
            # 创建测试文件
            test_dir = create_test_files(temp_base, size)
            
            # 测试原始方法
            original_result = test_original_method(test_dir, size)
            
            # 测试队列方法
            queued_result = test_queued_method(test_dir, size)
            
            if original_result and queued_result:
                # 计算性能提升
                speedup = original_result['total_time'] / queued_result['total_time']
                
                print(f"\n📈 性能对比结果 (文件数: {size})")
                print(f"   原始方法总时间: {original_result['total_time']:.2f}秒")
                print(f"   队列方法总时间: {queued_result['total_time']:.2f}秒")
                print(f"   性能提升: {speedup:.2f}倍")
                print(f"   原始方法速度: {original_result['avg_speed']:.1f} 文件/秒")
                print(f"   队列方法速度: {queued_result['index_speed']:.1f} 文件/秒")
                
                results.append({
                    'file_count': size,
                    'original': original_result,
                    'queued': queued_result,
                    'speedup': speedup
                })
            
            # 清理测试目录
            shutil.rmtree(test_dir)
    
    finally:
        # 清理临时目录
        shutil.rmtree(temp_base)
    
    # 显示总结
    show_summary(results)


def show_summary(results):
    """显示性能对比总结"""
    if not results:
        print("❌ 没有有效的测试结果")
        return
    
    print("\n" + "=" * 70)
    print("📊 性能对比总结")
    print("=" * 70)
    
    print(f"{'文件数量':<10} {'原始方法(秒)':<15} {'队列方法(秒)':<15} {'性能提升':<10}")
    print("-" * 70)
    
    for result in results:
        print(f"{result['file_count']:<10} "
              f"{result['original']['total_time']:<15.2f} "
              f"{result['queued']['total_time']:<15.2f} "
              f"{result['speedup']:<10.2f}x")
    
    # 计算平均性能提升
    avg_speedup = sum(r['speedup'] for r in results) / len(results)
    print("-" * 70)
    print(f"平均性能提升: {avg_speedup:.2f}倍")
    
    print("\n💡 结论:")
    if avg_speedup > 1.5:
        print(f"   队列模式相比原始方法有显著性能提升 ({avg_speedup:.1f}倍)")
    elif avg_speedup > 1.0:
        print(f"   队列模式相比原始方法有一定性能提升 ({avg_speedup:.1f}倍)")
    else:
        print("   两种方法性能相近，建议根据具体场景选择")


if __name__ == "__main__":
    print("选择测试模式:")
    print("1. 完整性能对比测试")
    print("2. 自定义文件数量测试")
    
    choice = input("请选择 (1-2): ").strip()
    
    if choice == "1":
        compare_performance()
    elif choice == "2":
        file_count = int(input("请输入测试文件数量: "))
        temp_base = tempfile.mkdtemp(prefix="custom_test_")
        
        try:
            test_dir = create_test_files(temp_base, file_count)
            original_result = test_original_method(test_dir, file_count)
            queued_result = test_queued_method(test_dir, file_count)
            
            if original_result and queued_result:
                speedup = original_result['total_time'] / queued_result['total_time']
                print(f"\n📈 性能对比: 队列方法比原始方法快 {speedup:.2f}倍")
        finally:
            shutil.rmtree(temp_base)
    else:
        print("无效选择")